# -*- coding: utf-8 -*-
#voc yolo格式转换，yolo训练集制作
import os
import random
import shutil
import cv2
from lxml import etree
from xml.dom.minidom import Document
def voc2yolo(voc_img_dir, voc_label_dir, yolo_img_dir, yolo_label_dir,classes,clear_flag = False):#参考#https://www.cnblogs.com/6S6B/p/16544237.html
    xmls = os.listdir(voc_label_dir)
    xmls = [x for x in xmls if x.endswith('.xml')]
    if yolo_img_dir is not None:
        if not os.path.exists(yolo_img_dir):
            os.mkdir(yolo_img_dir)
    if not os.path.exists(yolo_label_dir):
        os.mkdir(yolo_label_dir)
    all_xmls = len(xmls)
    for idx, one_xml in enumerate(xmls):
        xl = etree.parse(os.path.join(voc_label_dir, one_xml))
        root = xl.getroot()
        objects = root.findall('object')
        img_size = root.find('size')
        img_w = 0
        img_h = 0
        if img_size:
            img_width = img_size.find('width')
            if img_width is not None:
                img_w = int(img_width.text)
            img_height = img_size.find('height')
            if img_height is not None:
                img_h = int(img_height.text)
        label_lines = []
        for ob in objects:
            one_annotation = {}
            label = ob.find('name').text
            one_annotation['tag'] = label
            one_annotation['flag'] = False
            bbox = ob.find('bndbox')
            xmin = int(bbox.find('xmin').text)
            ymin = int(bbox.find('ymin').text)
            xmax = int(bbox.find('xmax').text)
            ymax = int(bbox.find('ymax').text)
            if img_w == 0 or img_h == 0:
                img = cv2.imread(os.path.join(voc_img_dir, one_xml.replace('.xml', '.jpg')))
                img_h, img_w = img.shape[:2]
            bbox_w = (xmax - xmin) / img_w
            bbox_h = (ymax - ymin) / img_h
            bbox_cx = (xmin + xmax) / 2 / img_w
            bbox_cy = (ymin + ymax) / 2 / img_h
            try:
                bbox_label = classes.index(label)#class_num[label]
                label_lines.append(f'{bbox_label} {bbox_cx} {bbox_cy} {bbox_w} {bbox_h}' + '\n')
            except Exception as e:
                print("not find number label in class_num ", e, one_xml)
                label_lines = []
                break
        if len(label_lines):
            with open(os.path.join(yolo_label_dir, one_xml.replace('.xml', '.txt')), 'w') as fp:
                fp.writelines(label_lines)
            if yolo_img_dir is not None:
                shutil.copy(os.path.join(voc_img_dir, one_xml.replace('.xml', '.jpg')),
                            os.path.join(yolo_img_dir))
        print(f"processing: {idx}/{all_xmls}")

        if clear_flag == True:
            shutil.rmtree(voc_img_dir)
            print("dir ",voc_img_dir," cleared")
            os.mkdir(voc_img_dir)
            print("dir ",voc_img_dir," created")



"""此函数用于将yolo格式txt标注文件转换为voc格式xml标注文件"""
def yolo2voc(yolo_img_dir, yolo_label_dir,voc_img_dir, voc_label_dir,classes, clear_flag):
    files = os.listdir(yolo_label_dir)
    print(len(files))
    
    for i, name in enumerate(files):
        xmlBuilder = Document()
        annotation = xmlBuilder.createElement("annotation")  # 创建annotation标签
        xmlBuilder.appendChild(annotation)
        txtFile = open(os.path.join(yolo_label_dir , name))
        txtList = txtFile.readlines()
        img = cv2.imread(os.path.join(yolo_img_dir ,name[0:-4] + ".jpg"))
        Pheight, Pwidth, Pdepth = img.shape

        folder = xmlBuilder.createElement("folder")  # folder标签
        foldercontent = xmlBuilder.createTextNode("driving_annotation_dataset")
        folder.appendChild(foldercontent)
        annotation.appendChild(folder)  # folder标签结束

        filename = xmlBuilder.createElement("filename")  # filename标签
        filenamecontent = xmlBuilder.createTextNode(name[0:-4] + ".jpg")
        filename.appendChild(filenamecontent)
        annotation.appendChild(filename)  # filename标签结束

        size = xmlBuilder.createElement("size")  # size标签
        width = xmlBuilder.createElement("width")  # size子标签width
        widthcontent = xmlBuilder.createTextNode(str(Pwidth))
        width.appendChild(widthcontent)
        size.appendChild(width)  # size子标签width结束

        height = xmlBuilder.createElement("height")  # size子标签height
        heightcontent = xmlBuilder.createTextNode(str(Pheight))
        height.appendChild(heightcontent)
        size.appendChild(height)  # size子标签height结束

        depth = xmlBuilder.createElement("depth")  # size子标签depth
        depthcontent = xmlBuilder.createTextNode(str(Pdepth))
        depth.appendChild(depthcontent)
        size.appendChild(depth)  # size子标签depth结束

        annotation.appendChild(size)  # size标签结束

        for j in txtList:
            oneline = j.strip().split(" ")
            object = xmlBuilder.createElement("object")  # object 标签
            picname = xmlBuilder.createElement("name")  # name标签
            namecontent = xmlBuilder.createTextNode(classes[int(oneline[0])]) #xmlBuilder.createTextNode(dic[oneline[0]])
            picname.appendChild(namecontent)
            object.appendChild(picname)  # name标签结束

            pose = xmlBuilder.createElement("pose")  # pose标签
            posecontent = xmlBuilder.createTextNode("Unspecified")
            pose.appendChild(posecontent)
            object.appendChild(pose)  # pose标签结束

            truncated = xmlBuilder.createElement("truncated")  # truncated标签
            truncatedContent = xmlBuilder.createTextNode("0")
            truncated.appendChild(truncatedContent)
            object.appendChild(truncated)  # truncated标签结束

            difficult = xmlBuilder.createElement("difficult")  # difficult标签
            difficultcontent = xmlBuilder.createTextNode("0")
            difficult.appendChild(difficultcontent)
            object.appendChild(difficult)  # difficult标签结束

            bndbox = xmlBuilder.createElement("bndbox")  # bndbox标签
            xmin = xmlBuilder.createElement("xmin")  # xmin标签
            mathData = int(((float(oneline[1])) * Pwidth + 1) - (float(oneline[3])) * 0.5 * Pwidth)
            xminContent = xmlBuilder.createTextNode(str(mathData))
            xmin.appendChild(xminContent)
            bndbox.appendChild(xmin)  # xmin标签结束

            ymin = xmlBuilder.createElement("ymin")  # ymin标签
            mathData = int(((float(oneline[2])) * Pheight + 1) - (float(oneline[4])) * 0.5 * Pheight)
            yminContent = xmlBuilder.createTextNode(str(mathData))
            ymin.appendChild(yminContent)
            bndbox.appendChild(ymin)  # ymin标签结束

            xmax = xmlBuilder.createElement("xmax")  # xmax标签
            # 注意这里对于渗水的标注是会超出界限的，因为渗水是整幅图
            mathData = int(((float(oneline[1])) * Pwidth + 1) + (float(oneline[3])) * 0.5 * Pwidth)
            xmaxContent = xmlBuilder.createTextNode(str(mathData))
            xmax.appendChild(xmaxContent)
            bndbox.appendChild(xmax)  # xmax标签结束

            ymax = xmlBuilder.createElement("ymax")  # ymax标签
            mathData = int(((float(oneline[2])) * Pheight + 1) + (float(oneline[4])) * 0.5 * Pheight)
            ymaxContent = xmlBuilder.createTextNode(str(mathData))
            ymax.appendChild(ymaxContent)
            bndbox.appendChild(ymax)  # ymax标签结束

            object.appendChild(bndbox)  # bndbox标签结束

            annotation.appendChild(object)  # object标签结束

        f = open(os.path.join(voc_label_dir ,name[0:-4] + ".xml"), 'w')
        xmlBuilder.writexml(f, indent='\t', newl='\n', addindent='\t', encoding='utf-8')
        f.close()
    
    for filename in os.listdir(yolo_img_dir):
        imgfile = os.path.join(yolo_img_dir,filename)
        shutil.copy(imgfile,voc_img_dir)

    if clear_flag == True:
        shutil.rmtree(yolo_img_dir)
        os.makedirs(yolo_img_dir)



#根据voc_label 即xml文件，按比例生成 txt文件
# 参考 https://blog.csdn.net/qq_40716944/article/details/128648001    YOLOv8训练自己的数据集(超详细)
def get_split_files(voc_label_dir, split_file_dir,ratios=(0.6,0.2,0.2)):

    #TODO 后续追加容错
    #0 生成目标文件夹，清空原内容（将原来的文件改名为backup 时间戳，以避免清错找不不回）

    #1 声明存放输出结果的各文件名和文件句柄
    #1.1 文件名
    train_file_name = 'train.txt'
    val_file_name = 'val.txt'
    test_file_name = 'test.txt'
    train_val_file_name = 'train_val.txt'
    #1.2 文件句柄
    train_val_file = open(os.path.join(split_file_dir, train_val_file_name),'w')
    test_file = open(os.path.join(split_file_dir, test_file_name), 'w')
    train_file = open(os.path.join(split_file_dir, train_file_name), 'w')
    val_file = open(os.path.join(split_file_dir, val_file_name), 'w')

    #2 遍历voc_label_dir ,按指定比例分配 ，写各个txt文件内容，每一行对应一个文件名（不含后缀）
    xml_file_names = os.listdir(voc_label_dir)
    file_nums = len(xml_file_names)
    total_idxs = range(file_nums)
    train_val_percent = ratios[0] + ratios[1]
    train_percent = ratios[0]
    val_percent = ratios[1]
    test_percent = ratios[2]

    #2.1 计算索引
    train_val_idxs = random.sample(total_idxs, int(file_nums * train_val_percent))
    train_idxs = random.sample(train_val_idxs, int(file_nums * train_percent))
    if len(train_idxs) < 1:
        print("no train samples!")

    # 写文件
    for idx in total_idxs:
        filename_without_extention = xml_file_names[idx][:-4] + "\n"   # 根据索引找到文件名，并去后缀   aaaa.xml ->  aaaa
        if idx in train_val_idxs:
            #写train_val
            train_val_file.write(filename_without_extention)
            if idx in train_idxs:
                #写train
                train_file.write(filename_without_extention)
            else:
                #写val
                val_file.write(filename_without_extention)
        else:
            #写test
            test_file.write(filename_without_extention)

    #3 关闭文件
    train_val_file.close()
    train_file.close()
    val_file.close()
    test_file.close()


if __name__ == "__main__":

    # #1 切分数据集
    # voc_label_dir =  ""
    # split_file_dir = ""
    # ratios = (0.6, 0.2, 0.2)
    # get_split_files(voc_label_dir,split_file_dir,ratios)

    # #2 voc转换为yolo格式
    voc_img_dir = "/home/cxz/yolo/data/images/val"
    voc_label_dir = "/home/cxz/yolo/data/labels/val"
    yolo_img_dir = "/home/cxz/downloads/val1"
    yolo_label_dir = "/home/cxz/downloads/val2"
    classes = ["single","double","single screw","double screw","stain","discard"]
    clear_flag = False #即将图片copy到yolo文件夹后，清除原始图片数据
    voc2yolo(voc_img_dir, voc_label_dir, yolo_img_dir, yolo_label_dir, classes, clear_flag)



    # #3 yolo转为voc格式
    # yolo2voc(yolo_img_dir, yolo_label_dir,voc_img_dir, voc_label_dir,classes, clear_flag)
